//-------------------------------------------------------------
// Purpose: This program simulates a connect-four game.
//          Moves are based on the relative values of      
//          different length sequences of B or W pieces.
// Author:  John Gauch
//-------------------------------------------------------------
#include <cmath>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <string>
using namespace std;

const int NUM_WEIGHTS = 7;
const int NUM_SETS = 300;
float weights[NUM_SETS][NUM_WEIGHTS] =
   { {10,20,40,80,20,40,80} };
int win_count[NUM_SETS];
char winner[NUM_SETS][NUM_SETS];

class Game
{
public:
   // Public constants
   static const char EMPTY = ' ';
   static const char BLACK = 'B';
   static const char WHITE = 'W';

   // Public methods
   Game();
   void set_weights(char player, float weights[NUM_WEIGHTS]);
   char play_game(char first_player, bool show_output);

private:
   // Private constants and variables
   static const int LEN = 4;
   static const int ROWS = 6;
   static const int COLS = 7;
   int black_weight[LEN][LEN];
   int white_weight[LEN][LEN];
   char board[ROWS][COLS];
   int filled[COLS];
   int move_count;

   // Private methods
   void clear_board();
   bool check_move(int col);
   int score_sequence(int row, int col, int drow, int dcol, char player);
   int score_move(int col, char player);
   void find_move(int &col, char player);
   void make_move(int col, char player);
   void show_board();
   bool check_win(char player);
};


//-------------------------------------------------------------
Game::Game()
{
   // Set all weights to 1
   for (int i = 0; i < LEN; i++)
      for (int j = 0; j < LEN; j++)
      {
         black_weight[i][j] = 1;
         white_weight[i][j] = 1;
      }

   clear_board();
}

//-------------------------------------------------------------
void Game::set_weights(char player, float weights[NUM_WEIGHTS])
{
   // Set all black weights 
   if (player == BLACK)
   {
      int index = 0;
      for (int j = 0; j < LEN; j++)
         black_weight[0][j] = weights[index++];
      for (int i = 1; i < LEN; i++)
         black_weight[i][0] = weights[index++];
      for (int i = 1; i < LEN; i++)
      for (int j = 1; j < LEN; j++)
         black_weight[i][j] = -1;
   }

   // Set all white weights 
   if (player == WHITE)
   {
      int index = 0;
      for (int j = 0; j < LEN; j++)
         white_weight[0][j] = weights[index++];
      for (int i = 1; i < LEN; i++)
         white_weight[i][0] = weights[index++];
      for (int i = 1; i < LEN; i++)
      for (int j = 1; j < LEN; j++)
         white_weight[i][j] = -1;
   }
}

//-------------------------------------------------------------
void Game::clear_board()
{
   // loop over all rows and cols
   for (int c = 0; c < COLS; c++)
   for (int r = 0; r < ROWS; r++)
      board[r][c] = EMPTY;
   for (int c = 0; c < COLS; c++)
      filled[c] = 0;
}

//-------------------------------------------------------------
bool Game::check_move(int col)
{
   // check location is valid
   return (col >= 0) && (col < COLS) && (filled[col] < ROWS);
}

//-------------------------------------------------------------
int Game::score_sequence(int row, int col, int drow, int dcol, char player)
{
   // count black/white/total pieces
   int black_count = 0;
   int white_count = 0;
   int total_count = 0;
   for (int step = 0; step < LEN; step++)
   {
      int r = row + step * drow;
      int c = col + step * dcol;
      if ((r >= 0) && (r < ROWS) && (c >= 0) && (c < COLS))
      {
         if (board[r][c] == BLACK) black_count++;
         if (board[r][c] == WHITE) white_count++;
         total_count++;
      }
   }

   // calculate score
   if ((total_count == LEN) && (player == WHITE))
      return white_weight[black_count][white_count];
   else if ((total_count == LEN) && (player == BLACK))
      return black_weight[white_count][black_count];
   else
      return -1;
}

//-------------------------------------------------------------
int Game::score_move(int col, char player)
{
   int score = random() % 10;
   if (check_move(col))
   {
      // check all offsets
      int row = filled[col];
      for (int step = 0; step < LEN; step++)
      {
         // check all directions
         score += score_sequence(row-step, col, 1, 0, player); 
         score += score_sequence(row, col-step, 0, 1, player); 
         score += score_sequence(row-step, col-step,  1, 1, player); 
         score += score_sequence(row+step, col-step, -1, 1, player); 
      }
   }
   return score; 
}

//-------------------------------------------------------------
void Game::find_move(int &col, char player)
{
   // loop over all starting positions
   col = 0;
   int best_score = -1000;
   for (int c = 0; c < COLS; c++)
   if (check_move(c))
   {
      // calculate best score
      int score = score_move(c, player);
      if (score > best_score)
      {
         col = c;
         best_score = score;
      }
   }
}

//-------------------------------------------------------------
void Game::make_move(int col, char player)
{
   // make move
   if (check_move(col))
   {
      int row = filled[col]++;
      board[row][col] = player; 
      move_count++;
   }
}

//-------------------------------------------------------------
void Game::show_board()
{
   // print all rows and cols
   cout << "\nMove: " << move_count << endl;
   for (int r = ROWS-1; r >= 0; r--)
   {
      cout << "+";
      for (int c = 0; c < COLS; c++)
         cout << "---+";
      cout << "\n";
      cout << "| ";
      for (int c = 0; c < COLS; c++)
         cout << board[r][c] << " | ";
      cout << "\n";
   }
   cout << "+";
   for (int c = 0; c < COLS; c++)
      cout << "---+";
   cout << "\n";
}

//-------------------------------------------------------------
bool Game::check_win(char player)
{
   // loop over all starting positions
   for (int c = 0; c < COLS; c++)
   for (int r = 0; r < ROWS; r++)
   if (board[r][c] == player)
   {
      // check row
      int count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r+d < ROWS) && (board[r+d][c] == player)) count++;
      if (count == LEN) return true;
    
      // check col
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((c+d < COLS) && (board[r][c+d] == player)) count++;
      if (count == LEN) return true;
    
      // check diagonal
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r+d < ROWS) && (c+d < COLS) && (board[r+d][c+d] == player)) count++;
      if (count == LEN) return true;
    
      // check diagonal
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r-d >= 0) && (c+d < COLS) && (board[r-d][c+d] == player)) count++;
      if (count == LEN) return true;
   }
   return false;
}

//-------------------------------------------------------------
char Game::play_game(char first_player, bool show_output)
{
   // Initialize game
   int col = 0;
   move_count = 0;
   clear_board();

   // Make first move
   if (first_player == BLACK)
   {
      find_move(col, BLACK);
      make_move(col, BLACK);
      if (show_output) show_board();
      if (check_win(BLACK)) return BLACK;
   }

   // Loop until someone wins game
   while (move_count < ROWS*COLS)
   {
      // White moves
      find_move(col, WHITE);
      make_move(col, WHITE);
      if (show_output) show_board();
      if (check_win(WHITE)) return WHITE;

      // Black moves
      find_move(col, BLACK);
      make_move(col, BLACK);
      if (show_output) show_board();
      if (check_win(BLACK)) return BLACK;
   }

   // Game was a tie
   return EMPTY;
}

//-------------------------------------------------------------
void calculate_weights(Game &game)
{
   // Loop over all sets of weights
   for (int i=0; i<NUM_SETS; i++)
   {
      // Play games with different weights
      win_count[i] = 0;
      game.set_weights(game.WHITE, weights[i]);
      for (int j=0; j<NUM_SETS; j++)
      {
         game.set_weights(game.BLACK, weights[j]);
         winner[i][j] = game.play_game(game.WHITE, false);
         if (winner[i][j] == game.WHITE) win_count[i]++;
      }

      // Print results
      cout << i << " " << win_count[i] << " ";
      for (int w=0; w<NUM_WEIGHTS; w++)
         cout << weights[i][w] << " ";
      cout << endl;
   }
}

//-------------------------------------------------------------
void improve_weights(Game &game)
{
   // Perform genetic algorithm to improve weights
   for (int count=0; count<NUM_SETS; count++)
   {
   // Find worst set of weights
   int worst=0;
   for (int i=0; i<NUM_SETS; i++)
      if (win_count[i] < win_count[worst])
         worst = i;
   
   // Generate new weights
   int rand1 = random() % NUM_SETS; 
   int rand2 = random() % NUM_SETS; 
   for (int w=0; w<NUM_WEIGHTS; w++)
   {
      int vote = random() % 100;
      if (vote < 33)
         weights[worst][w] = weights[rand1][w];   
      else if (vote < 66)
         weights[worst][w] = weights[rand2][w];   
      else
         weights[worst][w] = (weights[rand1][w] + weights[rand2][w])/2;   
   }

   // Play game again
   win_count[worst] = 0;
   game.set_weights(game.WHITE, weights[worst]);
   for (int j=0; j<NUM_SETS; j++)
   {
      game.set_weights(game.BLACK, weights[j]);
      winner[worst][j] = game.play_game(game.WHITE, false);
      if (winner[worst][j] == game.WHITE) win_count[worst]++;
   }
  
   // Print results
   cout << worst << " " << win_count[worst] << " ";
   for (int w=0; w<NUM_WEIGHTS; w++)
      cout << weights[worst][w] << " ";
   cout << endl;
   }

   // Print results
   for (int i=0; i<NUM_SETS; i++)
   {
      cout << i << " " << win_count[i] << " ";
      for (int w=0; w<NUM_WEIGHTS; w++)
         cout << weights[i][w] << " ";
      cout << endl;
   }
}

//-------------------------------------------------------------
int main(int argc, char *argv[])
{
   // Initialize random game weights
   Game game;
   srandom(time(NULL));
   for (int s=1; s<NUM_SETS; s++)
   for (int w=0; w<NUM_WEIGHTS; w++)
      weights[s][w] = random() % 100;

   // Optimize game weights
   for (int count=0; count<10; count++)
   {
      cout << "----------- " << count << " -----------\n"; 
      calculate_weights(game);
      improve_weights(game);
   }
}